#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <openssl/rsa.h>
#include <openssl/evp.h>
#include <openssl/objects.h>
#include <openssl/x509.h>
#include <openssl/x509v3.h>
#include <openssl/err.h>
#include <openssl/pem.h>
#include <openssl/ssl.h>

#include "postgres.h"
#include "executor/executor.h"
#include "fmgr.h"

#ifdef PG_MODULE_MAGIC
PG_MODULE_MAGIC;
#endif


#define __text_data(x) ((char *)VARDATA(x))
#define __text_length(x) (VARSIZE(x) - VARHDRSZ)

char *__text_to_pchar(text *arg, int *length) {
	int size = VARSIZE(arg) - VARHDRSZ;
	char *data = (char *) malloc(size + 1);
	memcpy(data, VARDATA(arg), size);
	data[size] = '\0';
	*length = size;
	return data;
}

#define __bytea_to_pchar(a,l) __text_to_pchar(a,l)

text *__pchar_to_text(char *arg, int length) {
	text *data = (text *) malloc(VARHDRSZ + length);
	SET_VARSIZE(data, VARHDRSZ + length);
	memcpy(VARDATA(data), arg, length);
	return data;
}

bytea *__pchar_to_bytea(char *arg, int length) {
	bytea *data = (bytea *) malloc(VARHDRSZ + length);
	SET_VARSIZE(data, VARHDRSZ + length);
	memcpy(VARDATA(data), arg, length);
	return data;
}

#define __text(pchar) __pchar_to_text(pchar, strlen(pchar))
#define __empty_text __pchar_to_text("", 0)
#define __empty_bytea __pchar_to_bytea("", 0)



char *__base64_decode(const char *input, int *length)
{
	BIO *b64, *bmem;

	char *buffer = (char *)malloc(*length);
	memset(buffer, 0, *length);

	b64 = BIO_new(BIO_f_base64());
	if (strchr(input, '\n') == NULL)
		BIO_set_flags(b64, BIO_FLAGS_BASE64_NO_NL);
	bmem = BIO_new_mem_buf((void *)input, *length);
	bmem = BIO_push(b64, bmem);

	*length = BIO_read(bmem, buffer, *length);

	BIO_free_all(bmem);

	return buffer;
}

char *__base64_encode(const char *input, int *length) {

	BIO *bmem, *b64;
	BUF_MEM *bptr;
	char *buff;

	b64 = BIO_new(BIO_f_base64());
	bmem = BIO_new(BIO_s_mem());
	b64 = BIO_push(b64, bmem);
	BIO_write(b64, input, *length);
	BIO_flush(b64);
	BIO_get_mem_ptr(b64, &bptr);

	buff = (char *)malloc(bptr->length);
	memcpy(buff, bptr->data, bptr->length - 1);
	buff[bptr->length - 1] = 0;
	*length = bptr->length - 1;
	BIO_free_all(b64);

	return buff;
}




#define __OBJ_T_PRIVATE_KEY 1
#define __OBJ_T_PUBLIC_KEY 2
#define __OBJ_T_X509_CERT 3
#define __OBJ_T_X509_REQ 4

text *__openssl_object_to_text(void *obj, int obj_type) {
	text *rv;
	char *buff;
	BUF_MEM *bptr;
	int write_rc = 0;
	BIO *bmem = BIO_new(BIO_s_mem());
	if (!bmem)
		return NULL;
	switch (obj_type) {
	case __OBJ_T_PRIVATE_KEY: write_rc = PEM_write_bio_PrivateKey(bmem, (EVP_PKEY *) obj, NULL, NULL, 0, NULL, NULL); break;
	case __OBJ_T_PUBLIC_KEY: write_rc = PEM_write_bio_PUBKEY(bmem, (EVP_PKEY *) obj); break;
	case __OBJ_T_X509_CERT: write_rc = PEM_write_bio_X509(bmem, (X509 *) obj); break;
	case __OBJ_T_X509_REQ: write_rc = PEM_write_bio_X509_REQ(bmem, (X509_REQ *) obj); break;
	}
	if (!write_rc) {
		BIO_free(bmem);
		return NULL;
	}
	BIO_get_mem_ptr(bmem, &bptr);
	buff = (char *) malloc(bptr->length);
	memcpy(buff, bptr->data, bptr->length-1);
	buff[bptr->length - 1] = 0;
	rv = __pchar_to_text(buff, bptr->length - 1); // TODO optimize
	BIO_free(bmem);
	free(buff);
	return rv;
}

#define __private_key_to_text(key) __openssl_object_to_text((void *) key, __OBJ_T_PRIVATE_KEY)
#define __public_key_to_text(key) __openssl_object_to_text((void *) key, __OBJ_T_PUBLIC_KEY)
#define __x509_cert_to_text(cert) __openssl_object_to_text((void *) cert, __OBJ_T_X509_CERT)
#define __x509_req_to_text(cert) __openssl_object_to_text((void *) req, __OBJ_T_X509_REQ)




EVP_PKEY *__pchar_to_private_key(char *pem_key, int pem_key_length) {
	BIO *bmem = BIO_new_mem_buf((void *)pem_key, pem_key_length);
	EVP_PKEY *pkey = PEM_read_bio_PrivateKey(bmem, NULL, NULL, NULL);
	BIO_free(bmem);
	return pkey;
}

EVP_PKEY *__text_to_private_key(text *pem_key) {
	int length;
	char *src = __text_to_pchar(pem_key, &length);
	EVP_PKEY *pkey = __pchar_to_private_key(src, length);
	free(src);
	return pkey;
}

EVP_PKEY *__pchar_to_rsa_public_key(char *pem_key, int pem_key_length) {
	EVP_PKEY *pkey;
	BIO *bmem = BIO_new_mem_buf((void *)pem_key, pem_key_length);
	RSA *rsa = PEM_read_bio_RSA_PUBKEY(bmem, NULL, NULL, NULL);
	BIO_free(bmem);
	if (rsa == NULL)
		return NULL;
	pkey = EVP_PKEY_new();
	if (pkey == NULL)
		return NULL;
	EVP_PKEY_assign_RSA(pkey, rsa);
	return pkey;
}

EVP_PKEY *__text_to_rsa_public_key(text *pem_key) {
	int length;
	char *src = __text_to_pchar(pem_key, &length);
	EVP_PKEY *pkey = __pchar_to_rsa_public_key(src, length);
	free(src);
	return pkey;
}

EVP_PKEY *__pchar_to_dsa_public_key(char *pem_key, int pem_key_length) {
	EVP_PKEY *pkey;
	BIO *bmem = BIO_new_mem_buf((void *)pem_key, pem_key_length);
	DSA *dsa = PEM_read_bio_DSA_PUBKEY(bmem, NULL, NULL, NULL);
	BIO_free(bmem);
	if (dsa == NULL)
		return NULL;
	pkey = EVP_PKEY_new();
	if (pkey == NULL)
		return NULL;
	EVP_PKEY_assign_DSA(pkey, dsa);
	return pkey;
}

EVP_PKEY *__text_to_dsa_public_key(text *pem_key) {
	int length;
	char *src = __text_to_pchar(pem_key, &length);
	EVP_PKEY *pkey = __pchar_to_dsa_public_key(src, length);
	free(src);
	return pkey;
}

X509 *__pchar_to_x509_cert(char *pem_cert, int pem_cert_length) {
	BIO *bmem = BIO_new_mem_buf((void *)pem_cert, pem_cert_length);
	X509 *x509 = PEM_read_bio_X509(bmem, NULL, NULL, NULL);
	BIO_free(bmem);
	return x509;
}

X509 *__text_to_x509_cert(text *pem_cert) {
	int length;
	char *src = __text_to_pchar(pem_cert, &length);
	X509 *x509 = __pchar_to_x509_cert(src, length);
	free(src);
	return x509;
}

X509_REQ *__pchar_to_x509_req(char *pem_req, int pem_req_length) {
	BIO *bmem = BIO_new_mem_buf((void *)pem_req, pem_req_length);
	X509_REQ *x509_req = PEM_read_bio_X509_REQ(bmem, NULL, NULL, NULL);
	BIO_free(bmem);
	return x509_req;
}

X509_REQ *__text_to_x509_req(text *pem_req) {
	int length;
	char *src = __text_to_pchar(pem_req, &length);
	X509_REQ *x509_req = __pchar_to_x509_req(src, length);
	free(src);
	return x509_req;
}



text *base64_encode(bytea *_bin) {
	int length;
	char *bin;
	char *b64;
	text *result;
	bin = __bytea_to_pchar(_bin, &length);
	b64 = __base64_encode(bin, &length);
	free(bin);
	result = __pchar_to_text(b64, length);
	free(b64);
	return result;
}

bytea *base64_decode(text *_b64) {
	int length;
	char *bin;
	char *b64;
	bytea *result;
	b64 = __text_to_pchar(_b64, &length);
	bin = __base64_decode(b64, &length);
	free(b64);
	result = __pchar_to_bytea(bin, length);
	free(bin);
	return result;
}



text *gen_keys_pair(text *_alg, int32 num_bits) {

	RSA *rsa;
	DSA *dsa;
	EVP_PKEY *key;
	text *result;

	key = EVP_PKEY_new();
	if (key == NULL)
		return __empty_text;

	if (strncmp(__text_data(_alg), "rsa", __text_length(_alg)) == 0) {
		rsa = RSA_generate_key(num_bits, RSA_F4, NULL, NULL);
		if (rsa == NULL) {
			EVP_PKEY_free(key);
			return __empty_text;
		}
		if (!EVP_PKEY_assign_RSA(key, rsa)) {
			RSA_free(rsa);
			EVP_PKEY_free(key);
			return __empty_text;
		}
	} else if (strncmp(__text_data(_alg), "dsa", __text_length(_alg)) == 0) {
		dsa = DSA_generate_parameters(num_bits, NULL, 0, NULL, NULL, NULL, NULL);
		if (dsa == NULL || DSA_generate_key(dsa) != 1) {
			if (dsa)
				DSA_free(dsa);
			EVP_PKEY_free(key);
			return __empty_text;
		}
		if (!EVP_PKEY_assign_DSA(key, dsa)) {
			DSA_free(dsa);
			EVP_PKEY_free(key);
			return __empty_text;
		}
	} else {
		return __empty_text;
	}

	result = __private_key_to_text(key);
	EVP_PKEY_free(key);

	return (result ? result : __empty_text);
}

text *get_public_key(text *_private_key) {

	EVP_PKEY *key = __text_to_private_key(_private_key); /* pair of keys */
	if (key == NULL)
		return __empty_text;
	text *result = __public_key_to_text(key);
	EVP_PKEY_free(key);
	return (result ? result : __empty_text);
}



#define __RSA_OP_PRIVATE_ENCRYPT 0
#define __RSA_OP_PUBLIC_DECRYPT 1
#define __RSA_OP_PUBLIC_ENCRYPT 2
#define __RSA_OP_PRIVATE_DECRYPT 3

bytea *__rsa_crypt(text *_key, bytea *_data, int _op) {

	int data_length;
	int buff_length = -1;
	char *data;
	char *buff;
	EVP_PKEY *pk = NULL;
	RSA *rsa;
	bytea *result;

	switch (_op) {
	case __RSA_OP_PRIVATE_ENCRYPT: pk = __text_to_private_key(_key); break;
	case __RSA_OP_PUBLIC_DECRYPT: pk = __text_to_rsa_public_key(_key); break;
	case __RSA_OP_PUBLIC_ENCRYPT: pk = __text_to_rsa_public_key(_key); break;
	case __RSA_OP_PRIVATE_DECRYPT: pk = __text_to_private_key(_key); break;
	}
	if (!pk)
		return __empty_bytea;
	if (!(rsa = EVP_PKEY_get1_RSA(pk))) {
		EVP_PKEY_free(pk);
		return __empty_bytea;
	}
	EVP_PKEY_free(pk);
	buff = (char *) malloc(RSA_size(rsa));
	data = __bytea_to_pchar(_data, &data_length);
	
	switch (_op) {
	case __RSA_OP_PRIVATE_ENCRYPT: buff_length = RSA_private_encrypt(data_length, (unsigned char*)data, (unsigned char*)buff, rsa, RSA_PKCS1_PADDING); break;
	case __RSA_OP_PUBLIC_DECRYPT: buff_length = RSA_public_decrypt(data_length, (unsigned char*)data, (unsigned char*)buff, rsa, RSA_PKCS1_PADDING); break;
	case __RSA_OP_PUBLIC_ENCRYPT: buff_length = RSA_public_encrypt(data_length, (unsigned char*)data, (unsigned char*)buff, rsa, RSA_PKCS1_PADDING); break;
	case __RSA_OP_PRIVATE_DECRYPT: buff_length = RSA_private_decrypt(data_length, (unsigned char*)data, (unsigned char*)buff, rsa, RSA_PKCS1_PADDING); break;
	}
	RSA_free(rsa);
	free(data);

	if (buff_length == -1) {
		free(buff);
		return __empty_bytea;
	}

	result = __pchar_to_bytea(buff, buff_length);
	free(buff);

	return result;
}

bytea *rsa_private_encrypt(text *_private_key, bytea *_data) {
	return __rsa_crypt(_private_key, _data, __RSA_OP_PRIVATE_ENCRYPT);
}

bytea *rsa_public_decrypt(text *_public_key, bytea *_data) {
	return __rsa_crypt(_public_key, _data, __RSA_OP_PUBLIC_DECRYPT);
}

bytea *rsa_public_encrypt(text *_public_key, bytea *_data) {
	return __rsa_crypt(_public_key, _data, __RSA_OP_PUBLIC_ENCRYPT);
}

bytea *rsa_private_decrypt(text *_private_key, bytea *_data) {
	return __rsa_crypt(_private_key, _data, __RSA_OP_PRIVATE_DECRYPT);
}




bytea *__sign(EVP_MD *md, text *_key, bytea *_data) {

	int data_length;
	char *data;
	EVP_PKEY *pkey;
	EVP_MD_CTX md_ctx;
	text *result;
	int err;
	char *sig_buf;
	unsigned int sig_len;

	if ((pkey = __text_to_private_key(_key)) == NULL)
		return __empty_text;
	data = __bytea_to_pchar(_data, &data_length);

	sig_buf = (char *) malloc(EVP_PKEY_size(pkey));
	EVP_SignInit(&md_ctx, md);
	EVP_SignUpdate(&md_ctx, data, data_length);
	sig_len = sizeof(sig_buf);
	err = EVP_SignFinal(&md_ctx, (unsigned char*)sig_buf, &sig_len, pkey);
	EVP_PKEY_free(pkey);	
	free(data);

	if (err != 1)
		return __empty_text;

	result = __pchar_to_bytea(sig_buf, sig_len);
	free(sig_buf);

	return result;
}

EVP_MD *__get_ds_alg(text *_alg, int *is_rsa) {
	EVP_MD *md = NULL;
	*is_rsa = 1;
	if (0) {
	} else if (strncmp("rsa/md2", __text_data(_alg), __text_length(_alg)) == 0) {
		md = (EVP_MD *) EVP_md2();
	} else if (strncmp("rsa/md4", __text_data(_alg), __text_length(_alg)) == 0) {
		md = (EVP_MD *) EVP_md4();
	} else if (strncmp("rsa/md5", __text_data(_alg), __text_length(_alg)) == 0) {
		md = (EVP_MD *) EVP_md5();
	} else if (strncmp("rsa/sha1", __text_data(_alg), __text_length(_alg)) == 0) {
		md = (EVP_MD *) EVP_sha1();
	} else if (strncmp("rsa/sha224", __text_data(_alg), __text_length(_alg)) == 0) {
		md = (EVP_MD *) EVP_sha224();
	} else if (strncmp("rsa/sha256", __text_data(_alg), __text_length(_alg)) == 0) {
		md = (EVP_MD *) EVP_sha256();
	} else if (strncmp("rsa/sha384", __text_data(_alg), __text_length(_alg)) == 0) {
		md = (EVP_MD *) EVP_sha384();
	} else if (strncmp("rsa/sha512", __text_data(_alg), __text_length(_alg)) == 0) {
		md = (EVP_MD *) EVP_sha512();
	} else if (strncmp("dsa/dss1", __text_data(_alg), __text_length(_alg)) == 0) {
		md = (EVP_MD *) EVP_dss1();
		*is_rsa = 0;
	}
	return md;
}

bytea *sign(text *_alg, text *_private_key, bytea *_data) {
	int is_rsa;
	EVP_MD *md = __get_ds_alg(_alg, &is_rsa);
	if (!md)
		return __empty_text;
	return __sign(md, _private_key, _data);
}

int32 __verify(EVP_MD *md, text *arg_pub_part, bytea *arg_data, bytea *arg_signature, int is_rsa) {

	int pub_part_length;
	char *pub_part;
	X509 *cert = NULL;
	EVP_PKEY *pkey = NULL;
	char *sig_buf;
	char *data;
	int data_length;
	int sig_len;
	int err;
	EVP_MD_CTX md_ctx;

	pub_part = __text_to_pchar(arg_pub_part, &pub_part_length);
	if (strstr(pub_part, "-----BEGIN CERTIFICATE-----")) {
		if ((cert = __pchar_to_x509_cert(pub_part, pub_part_length)))
			pkey = X509_get_pubkey(cert);
	} else {
		if (is_rsa) {
			pkey = __pchar_to_rsa_public_key(pub_part, pub_part_length);
		} else {
			pkey = __pchar_to_dsa_public_key(pub_part, pub_part_length);
		}
	}
	free(pub_part);
	if (cert)
		X509_free(cert);
	if (pkey == NULL)
		return -1;

	data = __bytea_to_pchar(arg_data, &data_length);
	sig_buf = __bytea_to_pchar(arg_signature, &sig_len);

	EVP_VerifyInit(&md_ctx, md);
	EVP_VerifyUpdate(&md_ctx, data, data_length);
	err = EVP_VerifyFinal(&md_ctx, (unsigned char*)sig_buf, sig_len, pkey);
	EVP_PKEY_free(pkey);
	free(data);
	free(sig_buf);

	if (err != 1)
		return 0;

	return 1;
}

int32 verify(text *_alg, text *_public_key, bytea *_data, bytea *_signature) {
	int is_rsa;
	EVP_MD *md = __get_ds_alg(_alg, &is_rsa);
	if (!md)
		return 0;
	return __verify(md, _public_key, _data, _signature, is_rsa);
}




char *__make_block(int len, int fill, char *src, int src_len) {
        char *iv = (char *) malloc(len);
        if (src_len >= len) {
                memcpy(iv, src, len);
        } else {
                memset(&iv[src_len], fill, len - src_len);
                memcpy(iv, src, src_len);
        }
        return iv;
}

#define __OP_ENCRYPT 1
#define __OP_DECRYPT 0

#define __SALT_LENGHT 8

bytea *__sym_crypt(EVP_CIPHER *c, EVP_MD *md, bytea *_iv, text *_pw, int rounds, bytea *_salt, bytea *_data, int operation) {

	int p_length;
	int buff_length;
	int value_length;
	char *value;
	char *salt;
	char *iv;
	char *buff;
	unsigned char key[EVP_MAX_KEY_LENGTH];
	EVP_CIPHER_CTX ctx;
	bytea *result;

	value = __bytea_to_pchar(_salt, &value_length);
	salt = __make_block(__SALT_LENGHT, 0, value, value_length);
	free(value);

	value = __bytea_to_pchar(_iv, &value_length);
	iv = __make_block(EVP_MAX_IV_LENGTH, 0, value, value_length);
	free(value);

	value = __text_to_pchar(_pw, &value_length);
	EVP_BytesToKey(c, md, (unsigned char*)salt, (unsigned char*)value, value_length, rounds, key, (unsigned char*)iv);
	free(value);


	value = __bytea_to_pchar(_data, &value_length);
	buff = (char *) malloc(value_length + 1024*32);
	EVP_CIPHER_CTX_init(&ctx);
	EVP_CipherInit_ex(&ctx, c, NULL, key, (unsigned char*)iv, operation);
	EVP_CipherUpdate(&ctx, (unsigned char*)buff, &buff_length, (unsigned char*)value, value_length);
	EVP_CipherFinal_ex(&ctx, (unsigned char*)(buff + buff_length), &p_length); 
	EVP_CIPHER_CTX_cleanup(&ctx);
	free(value);
	free(iv);
	free(salt);

	buff_length += p_length;
	result = __pchar_to_bytea(buff, buff_length);
	free(buff);

	return result;
}

EVP_MD *__get_md_alg(text *_alg) {
	EVP_MD *md = NULL;
	if (0) {
	} else if (strncmp("md2", __text_data(_alg), __text_length(_alg)) == 0) {
		md = (EVP_MD *) EVP_md2();
	} else if (strncmp("md4", __text_data(_alg), __text_length(_alg)) == 0) {
		md = (EVP_MD *) EVP_md4();
	} else if (strncmp("md5", __text_data(_alg), __text_length(_alg)) == 0) {
		md = (EVP_MD *) EVP_md5();
	} else if (strncmp("sha1", __text_data(_alg), __text_length(_alg)) == 0) {
		md = (EVP_MD *) EVP_sha1();
	} else if (strncmp("sha224", __text_data(_alg), __text_length(_alg)) == 0) {
		md = (EVP_MD *) EVP_sha224();
	} else if (strncmp("sha256", __text_data(_alg), __text_length(_alg)) == 0) {
		md = (EVP_MD *) EVP_sha256();
	} else if (strncmp("sha384", __text_data(_alg), __text_length(_alg)) == 0) {
		md = (EVP_MD *) EVP_sha384();
	} else if (strncmp("sha512", __text_data(_alg), __text_length(_alg)) == 0) {
		md = (EVP_MD *) EVP_sha512();
	}
	return md;
}

EVP_CIPHER *__get_sym_alg(text *_alg) {
	EVP_CIPHER *c = NULL;
	if (0) {

	} else if (strncmp("des-cbc", __text_data(_alg), __text_length(_alg)) == 0) {
		c = (EVP_CIPHER *) EVP_des_cbc();
	} else if (strncmp("des-ecb", __text_data(_alg), __text_length(_alg)) == 0) {
		c = (EVP_CIPHER *) EVP_des_ecb();

	} else if (strncmp("bf-cbc", __text_data(_alg), __text_length(_alg)) == 0) {
		c = (EVP_CIPHER *) EVP_bf_cbc();
	} else if (strncmp("bf-ecb", __text_data(_alg), __text_length(_alg)) == 0) {
		c = (EVP_CIPHER *) EVP_bf_ecb();

	} else if (strncmp("idea-cbc", __text_data(_alg), __text_length(_alg)) == 0) {
		c = (EVP_CIPHER *) EVP_idea_cbc();
	} else if (strncmp("idea-ecb", __text_data(_alg), __text_length(_alg)) == 0) {
		c = (EVP_CIPHER *) EVP_idea_ecb();

	} else if (strncmp("aes128-cbc", __text_data(_alg), __text_length(_alg)) == 0) {
		c = (EVP_CIPHER *) EVP_aes_128_cbc();
	} else if (strncmp("aes128-ecb", __text_data(_alg), __text_length(_alg)) == 0) {
		c = (EVP_CIPHER *) EVP_aes_128_ecb();

	} else if (strncmp("aes192-cbc", __text_data(_alg), __text_length(_alg)) == 0) {
		c = (EVP_CIPHER *) EVP_aes_128_cbc();
	} else if (strncmp("aes192-ecb", __text_data(_alg), __text_length(_alg)) == 0) {
		c = (EVP_CIPHER *) EVP_aes_128_ecb();

	} else if (strncmp("aes256-cbc", __text_data(_alg), __text_length(_alg)) == 0) {
		c = (EVP_CIPHER *) EVP_aes_128_cbc();
	} else if (strncmp("aes256-ecb", __text_data(_alg), __text_length(_alg)) == 0) {
		c = (EVP_CIPHER *) EVP_aes_128_ecb();

	}
	return c;
}

bytea *encrypt(text *algorithm, text *pw, text *md_algorithm, int rounds, bytea *salt, bytea *iv, bytea *data) {
	EVP_CIPHER *c = __get_sym_alg(algorithm);
	EVP_MD *md = __get_md_alg(md_algorithm);
	if (!c || !md)
		return __empty_bytea;
	return __sym_crypt(c, md, iv, pw, (rounds <= 0 ? 1 : rounds), salt, data, __OP_ENCRYPT);
}

bytea *decrypt(text *algorithm, text *pw, text *md_algorithm, int rounds, bytea *salt, bytea *iv, bytea *data) {
	EVP_CIPHER *c = __get_sym_alg(algorithm);
	EVP_MD *md = __get_md_alg(md_algorithm);
	if (!c || !md)
		return __empty_bytea;
	return __sym_crypt(c, md, iv, pw, (rounds <= 0 ? 1 : rounds), salt, data, __OP_DECRYPT);
}

